import pymongo
import pandas as pd
from datetime import datetime

class DatabaseManager:
    def __init__(self, config):
        self.config = config
        self.client = pymongo.MongoClient(config['mongodb']['connection_string'])
        self.db = self.client[config['mongodb']['database']]
        
    def save_stock_data(self, stock_code, data):
        """保存股票数据"""
        collection = self.db['stock_data']
        data['timestamp'] = datetime.now()
        data['stock_code'] = stock_code
        collection.insert_one(data)
        
    def save_news_data(self, news_data):
        """保存新闻数据"""
        collection = self.db['news_data']
        for item in news_data:
            item['timestamp'] = datetime.now()
            collection.insert_one(item)
            
    def save_analysis_result(self, stock_code, analysis_type, result):
        """保存分析结果"""
        collection = self.db['analysis_results']
        document = {
            'stock_code': stock_code,
            'analysis_type': analysis_type,
            'result': result,
            'timestamp': datetime.now()
        }
        collection.insert_one(document)
        
    def save_strategy(self, stock_code, strategy_type, strategy):
        """保存策略"""
        collection = self.db['strategies']
        document = {
            'stock_code': stock_code,
            'strategy_type': strategy_type,
            'strategy': strategy,
            'timestamp': datetime.now()
        }
        collection.insert_one(document)
        
    def save_notification(self, user_id, notification_type, content, status):
        """保存通知记录"""
        collection = self.db['notifications']
        document = {
            'user_id': user_id,
            'notification_type': notification_type,
            'content': content,
            'status': status,
            'timestamp': datetime.now()
        }
        collection.insert_one(document)
        
    def get_stock_data(self, stock_code, start_date, end_date):
        """获取历史股票数据"""
        collection = self.db['stock_data']
        query = {
            'stock_code': stock_code,
            'timestamp': {
                '$gte': datetime.strptime(start_date, '%Y%m%d'),
                '$lte': datetime.strptime(end_date, '%Y%m%d')
            }
        }
        return list(collection.find(query))
        
    def get_latest_strategy(self, stock_code, strategy_type="short"):
        """获取最新的策略"""
        collection = self.db['strategies']
        result = collection.find_one(
            {"stock_code": stock_code, "strategy_type": strategy_type},
            sort=[("timestamp", -1)]
        )
        return result
        
    def get_latest_analysis(self, stock_code, analysis_type="summary"):
        """获取最新的分析结果"""
        collection = self.db['analysis_results']
        result = collection.find_one(
            {"stock_code": stock_code, "analysis_type": analysis_type},
            sort=[("timestamp", -1)]
        )
        return result
        
    def save_performance_metric(self, data):
        """保存性能指标"""
        data["timestamp"] = datetime.now()
        return self.db.performance_metrics.insert_one(data)
        
    def get_performance_metrics(self, limit=20):
        """获取性能指标"""
        return list(self.db.performance_metrics.find().sort("timestamp", -1).limit(limit))
        
    def get_latest_strategy(self, stock_code, strategy_type):
        """获取最新策略"""
        collection = self.db['strategies']
        return collection.find_one(
            {'stock_code': stock_code, 'strategy_type': strategy_type},
            sort=[('timestamp', -1)]
        ) 